#include "Ray.h"
#include "Vector3.h"
#include "miro.h"
#include "AASquare.h"
#include <math.h>
#include <iostream>
#include <fstream>
#include <cstdio>
#include <cstdlib>

static Vector3 randDir(Vector3 n, float &sintheta, float &costheta)
{
	float sin2 = ((float) rand()) / RAND_MAX;

	sintheta = sqrt(sin2);
	costheta = cos(asin(sintheta));
	float phi = 2 * PI * (((float) rand()) / RAND_MAX);

	float x = -sintheta * sin(phi);
	float y = sintheta * cos(phi);

	float cos2 = 1.0f - (sintheta * sintheta);
	float z = sqrt(cos2);

	Vector3 tangentRay;
	if (fabs(n.x) > fabs(n.y))
		tangentRay = Vector3(0, 1, 0);
	else
		tangentRay = Vector3(1, 0, 0);
	tangentRay = cross(n, tangentRay);
	tangentRay.normalize();

	return x * tangentRay + y * cross(n, tangentRay) + z * n;
}

static Vector3 randUDir(Vector3 n, float &sintheta, float &costheta)
{
	float x, y, z;

	do
	{
		x = ((float) rand()) / RAND_MAX;
		y = ((float) rand()) / RAND_MAX;
		z = ((float) rand()) / RAND_MAX;
	} while (x * x + y * y + z * z > 1);

	Vector3 dir(x, y, z);
	dir.normalize();

	Vector3 tangentRay;
	if (fabs(n.x) > fabs(n.y))
		tangentRay = Vector3(0, 1, 0);
	else
		tangentRay = Vector3(1, 0, 0);
	tangentRay = cross(n, tangentRay);
	tangentRay.normalize();

	Vector3 result = dir.x * tangentRay + dir.y * cross(n, tangentRay) + dir.z
			* n;

	costheta = dot(n, dir);
	sintheta = sin(acos(costheta));

	return result;
}

// Task 1: Path Tracing
float pathTraceScene(bool verbose = false)
{
	int printInterval = 10000;
	int outputInterval = 1000;
	int numRays = 2000000;
	int countHits = 0;
	AASquare light = AASquare(2, Vector3(0, 10, 0), YAXIS);
	AASquare mirror = AASquare(2, Vector3(5, 5, 0), XAXIS);

	float P = 100; // light power
	float A = 4; // light area
	float L = P / A / PI; // Radiance per ray - figure this out later
	double totalE = 0;
	double totalL = 0;
	double E;
	float sintheta, costheta;

	std::ofstream f("pt.txt");

	for (int i = 1; i <= numRays; i++)
	{
		Vector3 randRay = randDir(Vector3(0, 1, 0), sintheta, costheta);
		Ray r = Ray(Vector3(0), randRay);
		HitInfo hit;
		if (light.intersect(hit, r))
		{
			countHits++;
			totalE += L * PI;
			totalL += L;
		}
		else if (mirror.intersect(hit, r))
		{
			Vector3 reflection = r.d - 2 * (dot(r.d, hit.N)) * hit.N;
			reflection.normalize();
			Ray r2 = Ray(hit.P, reflection);
			if (light.intersect(hit, r2))
			{
				countHits++;
				totalE += L * PI;
				totalL += L;
			}
		}

		if(!verbose && i % (numRays/1000) == 0)
		{
			fprintf(stderr,"Completed: %.1f%%\r", ((double) i)/((double) numRays/100.));
		}
		else if (verbose && i % printInterval == 0)
		{
			E = totalE / (i + 1);
			printf("%d\t%.8f\n", i, E);
		}

		if(i % outputInterval == 0)
		{
			E = totalE / (i + 1);
			f << i << ", " << E << "\n";
		}
		
	}

	f.close();

	E = totalE / numRays;
	if (verbose)
	{
		printf("%d\t%.5f\n", numRays, E);
		printf("L = %.5f\n", totalL / numRays);
	}
	else
	{
		printf("Completed: 100.0%%\n");
	}

	return E;
}

// Task 2: Progressive Photon Mapping
float ppmScene(bool verbose = false)
{
	int printInterval = 1000000;
	int outputInterval = 1000000;
	double radius = 0.5;
	double radius2 = radius*radius;
	double minRadius = 0.01;
	double power = 0;
	unsigned long long hitPhotons = 0;
	unsigned long long totalPhotons = 0;
	unsigned long long maxPhotons = 200000000;
	float alpha = 0.7;
	double scale = 0;
	double photonPower = 100;
	double E, L;
	AASquare light = AASquare(2.0f, Vector3(0, 10, 0), YAXIS);
	AASquare mirror = AASquare(2.0f, Vector3(5, 5, 0), XAXIS);
	AASquare plane = AASquare(2.0f * radius, Vector3(0, 0, 0), YAXIS);
	float sintheta, costheta;
	std::ofstream f("ppm.txt");

	while (totalPhotons < maxPhotons)
	{
		int m = 0;
		int photons_per_iter = 2;
		for (int j = 0; j < photons_per_iter; j++)
		{
			float lx = 2 * ((float) rand()) / RAND_MAX - 1;
			float lz = 2 * ((float) rand()) / RAND_MAX - 1;

			// use cosine pdf
			Vector3 rayDir = randDir(Vector3(0, -1, 0), sintheta, costheta);
			Ray r = Ray(Vector3(lx, 10, lz), rayDir);
			HitInfo hit;

			totalPhotons += 1;
			if (plane.intersect(hit, r))
			{
				if (hit.P.length2() < radius2)
				{
					m += 1;
				}
			}
			else if (mirror.intersect(hit, r))
			{
				Vector3 reflection = r.d - 2 * (dot(r.d, hit.N)) * hit.N;
				reflection.normalize();
				Ray r2 = Ray(hit.P, reflection);
				if (plane.intersect(hit, r2))
				{
					if (hit.P.length2() < radius2)
					{
						m += 1;
					}
				}
			}
		}
		if (m > 0)
		{
			scale = (alpha * hitPhotons + alpha * m) / (alpha * hitPhotons + m);
			radius2 *= scale;
			hitPhotons += m;
			power = (power + m * photonPower) * scale; // scaled by PI - cancelled out BRDF
		}


		if(!verbose && totalPhotons % (maxPhotons/1000) == 0)
		{
			fprintf(stderr,"Completed: %.1f%%\r", ((double) totalPhotons)/((double) maxPhotons/100.));
		}
		if (verbose && totalPhotons % 1000000 == 0)
		{
			E = power / totalPhotons / (PI * radius2);
			printf("%lld\t%lld\t%.4f\t%.4f\n", totalPhotons, hitPhotons, radius, E);
		}

		if(totalPhotons % outputInterval == 0)
		{
			E = power / totalPhotons / (PI * radius2);
			f << totalPhotons << ", " << E << ", " << radius2 << ", " << alpha * hitPhotons << "\n";
		}
	}
	
	f.close();
	
	E = power / totalPhotons / (PI * radius2);
	L = power / totalPhotons / (PI * radius2);
	if (verbose)
	{
		printf("%lld\t%lld\t%.4f\t%.4f\n", totalPhotons, hitPhotons, radius2, E);
	}
	else
	{
		printf("Completed: 100.0%%\n");
	}

	return E;
}

float p1(float theta, float r)
{
	float A = 2*2; // area of light
	return r*r/cos(theta)/A;
}

float p2(float theta)
{
	return cos(theta)/PI;
}

float w1(float theta, float r)
{
	return p1(theta, r) / (p1(theta, r) + p2(theta));
}

float w2(float theta, float r)
{
	return p2(theta) / (p1(theta, r) + p2(theta));
}

// Task 3: Multiple Importance Sampling
float misScene(bool verbose = false)
{
	int numRays = 2000000;
	int printInterval = 10000;
	int outputInterval = 1000;
	float c1 = 0.5; // fraction of total samples per type
	float c2 = 0.5;

	int countHits = 0;
	AASquare light = AASquare(2, Vector3(0, 10, 0), YAXIS);
	AASquare mirror = AASquare(2, Vector3(5, 5, 0), XAXIS);

	float P = 100; // light power
	float A = 4; // light area
	float L = P / A / PI; // Radiance per ray - figure this out later
	double E = 0, E1 = 0, E2 = 0, E3 = 0;
	double E1u = 0, E2u = 0;
	float sintheta, costheta;

	std::ofstream f("mis.txt");

	unsigned long long N1 = 0;
	unsigned long long N2 = 0;
	for (unsigned long long i = 1; i <= 0.5 * numRays; i++)
	{
		// Direct light sampling
		float lx = 2 * ((float) rand()) / RAND_MAX - 1;
		float lz = 2 * ((float) rand()) / RAND_MAX - 1;

		float r = Vector3(lx, 10, lz).length();
		float theta = acos(10/r);

		N1++;
		// it always hits, since there's nothing in the way
		E1 += w1(theta, r) * L * A * cos(theta) * cos(theta) / r / r; // TODO: is L the sample value?
		E1u += L * A * cos(theta) * cos(theta) / r / r;

		// BRDF sampling
		Vector3 randRay = randDir(Vector3(0, 1, 0), sintheta, costheta);
		Ray ray = Ray(Vector3(0), randRay);
		HitInfo hit;

		N2++;
		// direct hit
		if (light.intersect(hit, ray))
		{
			E2 += PI * w2(asin(sintheta), hit.t) * L;
			E2u += PI * L;
		}

		// indirect hit
		else if (mirror.intersect(hit, ray))
		{
			Vector3 reflection = ray.d - 2 * (dot(ray.d, hit.N)) * hit.N;
			reflection.normalize();
			Ray r2 = Ray(hit.P, reflection);
			if (light.intersect(hit, r2))
			{
				E3 += PI * L; // add in all of it, since it isn't seen by method 1
			}
		}

		if(!verbose && 2*i % (numRays/1000) == 0)
		{
			fprintf(stderr,"Completed: %.1f%%\r", ((double) i)/((double) numRays/200.));
		}
		else if (verbose && 2*i % printInterval == 0)
		{
			E = E1/N1 + (E2 + E3) / N2;
			printf("%d\t%.8f\t%.6f\t%.6f\t%.6f\n", (int) (2*i), E, E1/N1, E2/N2, E3/N2);
		}

		/*if(i<=1000)
		{
			f << i << ", " << E1u/N1 << "\n";
		}*/
		if(2*i % outputInterval == 0)
		{
			E = E1/N1 + (E2 + E3) / N2;
			f << 2*i << ", " << E << ", " << E1u/N1 << ", " << E2u/N2 << ", " << E3/N2 << "\n";
		}
	}

	f.close();

	E = E1/N1 + (E2 + E3) / N2;
	if (verbose)
	{
		printf("E = %.5f\n", E);
	}
	else
	{
		printf("Completed: 100.0%%\n");
	}

	return E;
}

int main(int argc, char*argv[])
{
  	//float E_pt = pathTraceScene(false);
	//float E_ppm = ppmScene(false);
	float E_mis = misScene(false);

	return 0;
}

